#include "xparameters.h"
#include "xuartlite.h"
#include <stdio.h>
#include <stdlib.h>
#include <math.h>

// 减小内存使用
#define IMAGE_SIZE 784
#define NUM_CLASSES 10
#define HIDDEN_SIZE 64  // 减少隐藏层大小

// MNIST数据存储在Flash中的地址
#define MNIST_DATA_ADDR 0x10000000

// UART实例
static XUartLite UartLite;

// 初始化串口
void init_uart() {
    XUartLite_Initialize(&UartLite, XPAR_AXI_UARTLITE_0_DEVICE_ID);
}

// 改写print_image函数使用UART输出
void print_image(float *image) {
    char buf[32];
    XUartLite_Send(&UartLite, (u8*)"\r\nInput Image:\r\n", 15);
    for(int i = 0; i < 28; i++) {
        for(int j = 0; j < 28; j++) {
            char c = (image[i * 28 + j] > 0.5f) ? '#' : '.';
            XUartLite_Send(&UartLite, (u8*)&c, 1);
        }
        XUartLite_Send(&UartLite, (u8*)"\r\n", 2);
    }
}

// 主函数
int main() {
    init_uart();
    
    // 初始化神经网络
    Network *net = create_network();
    
    // 从预定义位置加载神经网络权重
    // ...load weights from flash memory...
    
    while(1) {
        // 等待用户输入测试图像编号
        char input[10];
        int idx = 0;
        XUartLite_Send(&UartLite, (u8*)"Enter test number: ", 18);
        
        // 读取输入
        char c;
        while((c = XUartLite_Recv(&UartLite)) != '\r') {
            input[idx++] = c;
            XUartLite_Send(&UartLite, (u8*)&c, 1);
        }
        input[idx] = 0;
        
        int test_num = atoi(input);
        // 处理测试图像
        // ...process test image...
    }
    
    return 0;
}
